import kuai
import kuai.kuaiext
from kuai.ff import *
from kuai.sim.simtool import *
import numpy

class Energy:
    def __init__(self, name = ""):
        self.name = name
        
    def energies(self, mol, ff, energies):
        raise NotImplementedError()
    
    def forces(self, mol, ff, energies, forces):
        raise NotImplementedError()
    
    def hessian(self, mol, ff, energies, forces, hessian):
        raise NotImplementedError()

    def rank(self):
        raise NotImplementedError()

class EnergySet(list, Energy):
    def __init__(self, name = ""):
        Energy.__init__(self, name)
        
    def energies(self, mol, ff, energies):
        start = 0
        for i in self:
            next = start + i.rank()
            i.energies(mol, ff, energies[start, next])
            start = next
        
    def forces(self, mol, ff, energies, forces):
        start = 0
        for i in self:
            next = start + i.rank()
            i.forces(mol, ff, energies[start, next], forces)
            start = next

    def hessian(self, mol, ff, energies, forces, hessian):
        start = 0
        for i in self:
            next = start + i.rank()
            i.hessian(mol, ff, energies[start, next], forces, hessian)
            start = next

    def rank(self):
        result = 0
        for i in self:
            result += i.rank()
        return result

class ForceFieldEnergy(Energy):
    def __init__(self, mol, ff, term, pairs, name):
        Energy.__init__(self, name)
        lists = []
        for i in pairs:
            try:
                para = ff[ForceFieldKey(term, i)]
                lists.append(mol.index(i[0])*3)
                lists.append(mol.index(i[1])*3)
                lists.append(para.index)
            except KeyError:
                pass
        self.index = numpy.array(lists, dtype=ATOM_INDEX)

    def rank(self):
        return 1

class BondEnergy(ForceFieldEnergy):
    def __init__(self, mol, ff, term, name = None):
        if name == None:
            name = "Bond Energy (%s)" % term
        pairs = kuai.inner.list_atom_pair12(mol)
        ForceFieldEnergy.__init__(self, mol, ff, term, pairs, name)

class EnergyBHARM(BondEnergy):
    def __init__(self, mol, ff, term, name = None):
        BondEnergy.__init__(self, mol, ff, term, name)
        
    def energy(self, mol, ff):
        return kuai.kuaiext.bond_energy_harm_0(mol.coords, ff.k, self.index)
    
    def gradient(self, mol, ff, result):
        return kuai.kuaiext.bond_energy_harm_1(mol.coords, ff.k, self.index, result)
    

class NonbondEnergy(ForceFieldEnergy):
    def __init__(self, mol, ff, term, name = None):
        if name == None:
            Energy.__init__(self, "Nonbond Energy (%s)" % term)
        else:
            Energy.__init__(self, name)
        
        nonbonds = kuai.inner.list_nonbonds(mol)
        lists = []
        for i in nonbonds:
            key = ForceFieldKey(term, i)
            if i in ff:
                lists.append(mol.index(i[0]))
                lists.append(mol.index(i[1]))
                lists.append(ff[key].index)
        self.index = numpy.array(lists, dtype=ATOM_INDEX)
    
class EnergyN12_6(NonbondEnergy):
    def __init__(self, mol, ff, term, name = None):
        EnergyN12_6.__init__(self, mol, ff, term, name)


